"""
模块名：hxpITT
功能：截取屏幕，调用 CnOCR 转为文字，格式化文本
作者：hxp@海小皮
版本： V1.1
依赖：
    DataStruct
    cnocr
    pyautogui
    numpy
    cv2
    re
"""
from .DataStruct import Dialog, ScreenType, EmotionType
from cnocr import CnOcr
import pyautogui
import numpy as np
import cv2
import re
import PlotDub.TextProcess as tp

# 宏定义常量
MAX_ROLE_NAME_LEN = 10  # 角色名最大长度
# TODO: 这个阈值可能要根据分辨率变化，目前先设置为定值
SAME_LINE_THRESHOLD = 5  # 同行检测阈值


class ScreenParam:
    """
    屏幕截图参数数据结构
    成员：
        screenWidth 屏幕宽度 默认为 1920
        screenHeight 屏幕高度 默认为 1080
        margin 默认 1080p 屏幕截图框与屏幕边界 margin 为 60 像素
        dialogWidth 对话框宽度 根据前面的值自动推算
        dialogHeight 对话框高度 默认 1080p 时为 200
    """
    def __init__(self, s_width=1920, s_height=1080, s_margin=60, d_height=200):
        """
        s_width  # 屏幕宽度 默认为 1920
        s_height  # 屏幕高度 默认为 1080
        s_margin  # 默认 1080p 屏幕截图框与屏幕边界 margin 为 60 像素
        d_height  # 对话框高度 默认 1080p 下为 200
        """
        self.screenWidth = s_width  # 屏幕宽度 默认为 1920
        self.screenHeight = s_height  # 屏幕高度 默认为 1080
        self.margin = s_margin  # 默认 1080p 屏幕截图框与屏幕边界 margin 为 60 像素
        self.dialogWidth = self.screenWidth - self.margin * 2  # 对话框宽度
        self.dialogHeight = d_height  # 对话框高度 默认 1080p 下为 200

    def __repr__(self):
        return f"【屏幕参数ScreenParam】(【宽度】={self.screenWidth}, 【高度】={self.screenHeight}, 【边距】={self.margin}, 【对话框宽度】={self.dialogWidth}, 【对话框高度】={self.dialogHeight})\n"


# 屏幕参数字典 key=ScreenType value=ScreenParam对象
SCREEN_RESOLUTION_PARAM_DICT = {
    ScreenType.S_1080P: ScreenParam(1920, 1080, 60, 200),
    ScreenType.S_2K: ScreenParam(2560, 1440, 80, 300),  # TODO: 2k 参数未经过测试
    ScreenType.S_4K: ScreenParam(3840, 2160, 120, 400)
}


class hxpITT:
    """
    Image To Text, ITT
    功能：截取屏幕，调用 CnOCR 转为文字，格式化文本
    注意：默认参数为 1080p 参数
    成员：
        __screenShotLeft 截图框左上顶点 x 坐标
        __screenShotTop 截图框左上顶点 y 坐标
        __screenShotWidth 截图框宽度
        __screenShotHeight 截图框高度
        __reader CnOcr 对象
    """
    def __init__(self, screen_param: ScreenParam = SCREEN_RESOLUTION_PARAM_DICT[ScreenType.S_1080P]):
        """
        功能：将屏幕参数转换为截图参数
        输入：
            screen_param ScreenParam 数据结构对象 默认值为 1080p 分辨率下的参数
        """
        self.__screenShotLeft = screen_param.margin  # 截图框左上顶点 x 坐标
        self.__screenShotTop = screen_param.screenHeight - (
                    screen_param.dialogHeight + screen_param.margin)  # 截图框左上顶点 y 坐标
        if (screen_param.margin + screen_param.dialogWidth) > screen_param.screenWidth:
            self.__screenShotWidth = screen_param.screenWidth - screen_param.margin * 2
        self.__screenShotWidth = screen_param.dialogWidth  # 截图框宽度
        self.__screenShotHeight = screen_param.dialogHeight  # 截图框高度
        self.__reader = CnOcr()

    def ResetShotParam(self, screen_param: ScreenParam):
        """
        功能：根据输入的屏幕分辨率参数，重新计算截图参数
        输入：
            screen_param: ScreenParam 屏幕分辨率参数
        """
        self.__screenShotLeft = screen_param.margin  # 截图框左上顶点 x 坐标
        self.__screenShotTop = screen_param.screenHeight - (
                screen_param.dialogHeight + screen_param.margin)  # 截图框左上顶点 y 坐标
        if (screen_param.margin + screen_param.dialogWidth) > screen_param.screenWidth:
            self.__screenShotWidth = screen_param.screenWidth - screen_param.margin * 2
        self.__screenShotWidth = screen_param.dialogWidth  # 截图框宽度
        self.__screenShotHeight = screen_param.dialogHeight  # 截图框高度

    def ScreenShot(self):
        """
        功能：根据设定的截图参数截取屏幕
        输出：截取的图片 image
        """
        # 捕获屏幕区域
        screenshot = pyautogui.screenshot(region=(self.__screenShotLeft, self.__screenShotTop,
                                                  self.__screenShotWidth, self.__screenShotHeight))
        image = np.array(screenshot)

        # 将BGR图像转换为RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # 返回结果
        return image

    def ImageToDialog(self, image):
        """
        功能：对话截图转文本
        输入：image 对话区域截图
        输出：Dialog 结构的识别对话信息，若 roleName 属性等于 ”None“ 表示识别无效
        """
        # 进行文字识别
        result = self.__reader.ocr(image)

        if not hasattr(self, "last_speach_content_list"):
            self.last_speach_content_list = []  # 初始化静态变量

        # 过滤出汉字和标点符号并按行分组
        filtered_results = []
        line_groups = []  # 坐标组，标注行位置
        text_groups = []  # 文本组，与坐标组对应
        for detection in result:
            # 获取文本
            text = detection['text']

            # 低置信度剔除
            confidence = detection['score']
            if confidence < 0.2:
                continue

            # 使用正则表达式提取汉字和标点符号
            text = re.sub("[^\u4e00-\u9fa5，。！？?,.!]", "", text)

            # 使用正则表达式检查是否只包含汉字和标点符号
            if text:
                filtered_results.append(text)
            else:
                continue

            # 获取文本框
            box = detection['position']

            # 获取文本框的中心点坐标
            center_x = (box[0][0] + box[1][0]) // 2
            center_y = (box[0][1] + box[2][1]) // 2

            # 判断是否存在已有组与该文本属于同一行
            added_to_group = False
            line_count = 0
            for group in line_groups:
                group_center_y = sum(point[1] for point in group) // len(group)
                if abs(center_y - group_center_y) < SAME_LINE_THRESHOLD:  # 设置一个阈值，表示同一行的文本
                    group.append((center_x, center_y))
                    text_groups[line_count].append(text)
                    added_to_group = True  # 标志该文本以被添加到对应的行
                    break
                line_count += 1

            # 如果不存在已有组，则创建新的组
            if not added_to_group:
                line_groups.append([(center_x, center_y)])
                text_groups.append([text])

        # 将识别内容格式化
        dialog = Dialog("", "", EmotionType.CALM, [""])
        speach_content = ""
        current_line = 0
        if len(text_groups) >= 2:  # 总行数大于 2 才有可能是有效对话
            for line in text_groups:
                # 判断第一行是否为名字
                if current_line == 0:
                    # 名字长度检测
                    if len(line[0]) < MAX_ROLE_NAME_LEN:
                        dialog.roleName = line[0]
                        current_line += 1
                        continue
                    else:
                        break  # 检测不到有效的名字直接跳过
                # 判断第二行是头衔还是对话
                elif current_line == 1 and len(text_groups) >= 3:
                    line_text = ""
                    for string in line:
                        line_text += string
                    # 行数 >= 3，但第二行字数小于 20，第二行为头衔
                    if len(line_text) < 20:
                        dialog.title = line_text
                    else:  # 第二行是对话，该角色没有头衔
                        # dialog.speachContent += line_text
                        speach_content += line_text
                    current_line += 1
                # 处理剩余行
                else:
                    for group in text_groups[current_line:]:  # 后面的组存放的是对话
                        for content in group:
                            # dialog.speachContent += content
                            speach_content += content
                    # 情感分析（目前仅通过 EmotionType 的标点符号分析方法进行分析）
                    dialog.emotion = EmotionType.Analysis(speach_content)

            # 对 speach_content 按照标点符号切分，填入 speachContent
            speach_content_list = tp.cut_with_punctuation(speach_content)
            # 解决 OCR 可能识别不到最后的标点符号的 bug
            if len(speach_content_list) > 0 and len(self.last_speach_content_list) > 0:
                if speach_content_list[-1] == self.last_speach_content_list[-1]:
                    if not tp.is_complete_sentence(speach_content_list[-1]):
                        speach_content_list[-1] += "。"
            self.last_speach_content_list = speach_content_list
            # 对 speach_content_list 中的短句进行合并
            speach_content_list = tp.merge_short_sentence(speach_content_list, 10)
            # 处理后的 speach_content_list 填入 speachContent
            dialog.speachContent = speach_content_list

        return dialog
